import sys

import secretflow as sf
import spu
from secretflow.data.vertical import read_csv
from secretflow.ml.boost.ss_xgb_v import Xgb
from secretflow.ml.linear.ss_sgd import SSRegression

print(sys.argv)
alice_party, alice_ip = sys.argv[1].split('-')
bob_party, bob_ip = sys.argv[2].split('-')
self_party, self_ip = sys.argv[3].split('-')
cluster_port = sys.argv[4]
spu_port = sys.argv[5]

# Check the version of your SecretFlow
print('The version of SecretFlow: {}'.format(sf.__version__))

# In case you have a running secretflow runtime already.
sf.shutdown()

# sf.init(['alice', 'bob', 'carol'], address='local')
address = f'{self_ip}:19010'
parties = {
    alice_party: {
        'address': f'{alice_ip}:{cluster_port}'
    },
    bob_party: {
        'address': f'{bob_ip}:{cluster_port}'
    }
}
cluster_config = {'parties': parties, 'self_party': self_party}
cdef = {
    'nodes': [
        {
            'party': alice_party,
            'address': f'{alice_ip}:{spu_port}'
        },
        {
            'party': bob_party,
            'address': f'{bob_ip}:{spu_port}'
        },
    ],
    'runtime_config': {
        'protocol': spu.spu_pb2.SEMI2K,
        'field': spu.spu_pb2.FM128,
        'sigmoid_mode': spu.spu_pb2.RuntimeConfig.SIGMOID_REAL,
    },
}

runtime_env = None
sf.init(
    address=address,
    cluster_config=cluster_config,
    runtime_env=runtime_env,
    namespace=None)

alice, bob = sf.PYU(alice_party), sf.PYU(bob_party)

spu = sf.SPU(cdef)

input_path = {
    alice: '/workspace/tee/data/breast_hetero_guest.csv',
    bob: '/workspace/tee/data/breast_hetero_host.csv'
}
output_path = {
    alice: f'/tmp/alice_psi_{cluster_port}.csv',
    bob: f'/tmp/bob_psi_{cluster_port}.csv'
}
spu.psi_csv('id', input_path, output_path, alice_party)

train_vdf = read_csv(filepath=output_path, drop_keys=['id'], keys=['id'])
print(train_vdf)
print(train_vdf.shape)
print(train_vdf.columns)

train_x = train_vdf.drop(columns=['y'])
train_y = train_vdf['y']

lr_model = SSRegression(spu)
lr_model.fit(
    x=train_x,
    y=train_y,
    epochs=5,
    learning_rate=0.01,
    batch_size=64,
    sig_type='t1',
    reg_type='logistic',
    penalty='l2',
    l2_norm=0.5,
)

xgb = Xgb(spu)
params = {
    'num_boost_round': 3,
    'max_depth': 5,
    'sketch_eps': 0.25,
    'objective': 'logistic',
    'reg_lambda': 0.2,
    'subsample': 1,
    'colsample_by_tree': 1,
    'base_score': 0.5,
}
xgb_model = xgb.train(params=params, dtrain=train_x, label=train_y)

with open(f'{cluster_port}_{spu_port}', 'w', encoding='utf8') as fw:
    fw.write('ok')
